/**
 * @brief 红黑树的实现
 * @version 1.1
 * @author 宋炜
 * @date 2023-6-7 ~ 2023-12-25 
 */
 
 // 2023-12-25 使用std::shared_ptr只能指针代替裸指针

#pragma once
#include <atomic>
#include <memory>
#include <list>
#include <iterator>
#include <iostream>

#if !defined( THREAD_SAFE )
#    define THREAD_SAFE 1
#endif

#if THREAD_SAFE
#    include <mutex>
#endif
namespace wheels
{
	template <typename keyType , typename dataType >
	class rbTree {
	public:
		// 颜色相关变量定义
		using color = bool;
		const static bool      RED = true;
		const static bool      BLACK   = false;
		using value_type = typename std::decay< typename std::remove_pointer< dataType > :: type >::type;
		using key_type = typename std::decay< typename std::remove_pointer< keyType > :: type >> type;
		struct stNode;
		using nodePtr_t = std::shared_ptr< stNode >;
		// 树节点定义
		struct stNode {
			key_type              m_key;
			value_type            m_val;
			color                 m_color;
			nodePtr_t             p_left;
			nodePtr_t             p_right;

			stNode(const key_type& k, const value_type& value , color red = BLACK ) : m_key( k ) , m_val( value ), m_color( red ), p_left(nullptr), p_right(nullptr) {}
			stNode( key_type&& k, value_type&& value , color red = BLACK ) : m_key( k ) , m_val( value ), m_color( red ), p_left(nullptr), p_right(nullptr) {}
			stNode(const stNode& b ):m_key( b.m_key ) , m_val( b.m_val ), m_color( b.m_color ), p_left(nullptr), p_right(nullptr){}

			stNode& operator=( const stNode& b ){
				m_key = b.m_key;
				m_val = b.m_val;
				m_color = b.m_color;
				return *this;
			}

			stNode& operator=( stNode&& b ){
				m_key = std::move( b.m_key );
				m_val = std::move( b.m_val );
				m_color = b.m_color;
				return *this;
			}
			const key_type& key() const { return m_key; }
		};

		using node_t = stNode;
		// 迭代器定义
        class stIterator__ : public std::iterator< std::forward_iterator_tag , stNode >{
			friend class rbTree;
		private:
            nodePtr_t               p_node__;
            std::list< nodePtr_t >   m_stack__;   // 迭代时候需要将数据进行压栈
		public:
            stIterator__( nodePtr_t node ):p_node__( node ){}
            stIterator__( const stIterator__& b ): p_node__( b.p_node__ ),m_stack__(b.m_stack__){}
			/**
			 * @brief 采用先序遍历迭代。
			 */
            stIterator__& operator++(){
				nodePtr_t top = nullptr;
                if( m_stack__.size() > 0 ){
                    top =m_stack__.back();
				}
				// 优先访问左子树
				if( ( top == nullptr ) ||                  // 允许空栈的时候访问第一个节点
				    // 如果访问左节点的时候如果发现了已经访问了右节点说明左节点已经被访问过
				    //                   不要访问已经访问过的左节点     不要访问已经访问过右节点的左节点   
                    ( top != nullptr &&  top->p_left != p_node__ &&  top->p_right != p_node__ &&   p_node__->p_left != nullptr ) ){
                    m_stack__.push_back( p_node__ );
                    p_node__ = p_node__->p_left;
					return *this;
				}
				// 访问右子树
                if( top != nullptr &&   top->p_right != p_node__ &&  p_node__->p_right != nullptr ){
                    m_stack__.push_back( p_node__ );
                    p_node__ = p_node__->p_right;
					return *this;
				}
			
                p_node__ = top;
                m_stack__.pop_back();
                if( m_stack__.size() > 0 ){ // 递归调用以弹出栈顶
					return operator++();
				}else{ // 满足end()条件
                    p_node__ = nullptr;
				}
			
				return *this;
			}
			
            stIterator__ operator++(int){
                stIterator__ tmp( *this );
				operator ++();
				return tmp;
			}
		
            bool operator==( const stIterator__& b ) const{ return p_node__ == b.p_node__; }
            bool operator!=( const stIterator__& b ) const{ return p_node__ != b.p_node__; }
		
            value_type& operator*() { return p_node__->m_val; }
            value_type* operator->() { return &p_node__->m_val; }
		};
		// 基本类型定义
        using iterator = stIterator__;
	private:
        nodePtr_t   p_root__;            // 根节点
        size_t      m_count__;           // 节点数量
#if THREAD_SAFE
        std::mutex  m_mutex__;
#endif
	public:
        rbTree() : p_root__(nullptr) , m_count__(0) {}
		
		virtual ~rbTree(){
#if THREAD_SAFE
            std::lock_guard< std::mutex > lck( m_mutex__ );
#endif
		}

		const keyType& key( iterator it ) const{
            return it.p_node__->m_key;
		}
		/**
		 * @brief 清理内存
		 */
		void clear(){
#if THREAD_SAFE
            std::lock_guard< std::mutex > lck( m_mutex__ );
#endif
            p_root__.reset();
            m_count__ = 0;
		}

        iterator begin(){ return iterator( p_root__ ); }

		iterator end(){ return iterator( nullptr ); }
		/**
		 * @brief 查找数据
		 * @param key[ I ], 数据键名称
		 * @return 找到返回数据迭代器，否则返回end()
		 */
		iterator find( const keyType& key ){
            nodePtr_t curr = p_root__;
#if THREAD_SAFE
            std::lock_guard< std::mutex > lck( m_mutex__ );
#endif			
			// 查找节点
			while (curr && curr->m_key != key ) {
				if (key < curr->m_key) {
					curr = curr->p_left;
				} else {
					curr = curr->p_right;
				}
			}

			return iterator( curr );
		}
		/**
		 * @brief 判断是否存在数据
		 * @param key[ I ],
		 */
		bool has( const keyType& key ){
            nodePtr_t curr = p_root__;
#if THREAD_SAFE
            std::lock_guard< std::mutex > lck( m_mutex__ );
#endif
			// 查找节点
			while (curr && curr->m_key != key ) {
				if (key < curr->m_key) {
					curr = curr->p_left;
				} else {
					curr = curr->p_right;
				}
			}

			return ( curr != nullptr );
		}
		/**
		 * @brief 取数据内容。重载了下标运算符，方便获取数据
		 */
		dataType& operator[]( const keyType& key ){
            nodePtr_t curr = p_root__;
#if THREAD_SAFE
            std::lock_guard< std::mutex > lck( m_mutex__ );
#endif
			// 查找节点
			while (curr && curr->m_key != key ) {
				if (key < curr->m_key) {
					curr = curr->p_left;
				} else {
					curr = curr->p_right;
				}
			}

			if( curr == nullptr ) return {};
		
			return curr->m_val;
		}
		/**
		 * @brief 插入节点
		 * @param key[ I ], 要插入的数据键
		 * @param val[ I ]， 要插入的数据的值
		 * @return 成功操作返回true,否则返回false
		 */
		bool insert(const keyType& key , const dataType& val) {
			bool ret = true;
			try{
#if THREAD_SAFE
                std::lock_guard< std::mutex > lck( m_mutex__ );
#endif				
                if (!p_root__) {
                    p_root__ = new stNode( key , val , BLACK );
					return true;
				}

                nodePtr_t curr = p_root__ ,  parent;

				while (curr) {
					parent = curr;

					if (key < curr->m_key) {
						curr = curr->p_left;
					} else if( key > curr->m_key ){
						curr = curr->p_right;
					}else{ // 说明数据已经存在
						std::cerr << "数据已经存在" << std::endl;
						return false;
					}
				}
				
				if ( key < parent->m_key) {
					parent->p_left.reset( stNode( key , val , RED ) );
                    m_count__ ++;
                        fix_insertion__(parent->p_left);
				} else {
					parent->p_right.reset( stNode( key , val , RED ) );
                    m_count__ ++;
                    fix_insertion__(parent->p_right );
				}
			}catch( std::bad_alloc& e ){
				std::cerr << e.what() << std::endl;
				ret = false;
			}

			return ret;
		}

		bool insert(const keyType& key , dataType&& val) {
			bool ret = true;
			try{
#if THREAD_SAFE
                std::lock_guard< std::mutex > lck( m_mutex__ );
#endif				
                if (!p_root__) {
                    p_root__.reset( stNode( key , val , BLACK ) );
					return true;
				}

                nodePtr_t curr = p_root__ , parent;

				while (curr) {
					parent = curr;

					if (key < curr->m_key) {
						curr = curr->p_left;
					} else if( key > curr->m_key ){
						curr = curr->p_right;
					}else{ // 说明数据已经存在
						std::cerr << "数据已经存在" << std::endl;
						return false;
					}
				}
				
				if ( key < parent->m_key) {
					parent->p_left = new stNode( key , val , RED );
                    m_count__ ++;
                        fix_insertion__(parent->p_left);
				} else {
					parent->p_right = new stNode( key , val , RED );
                    m_count__ ++;
                    fix_insertion__(parent->p_right );
				}
			}catch( std::bad_alloc& e ){
				std::cerr << e.what() << std::endl;
				ret = false;
			}

			return ret;
		}
		/**
		 * @brief 删除操作节点
		 */
		void erase(const keyType& key ) {
            nodePtr_t curr = p_root__;
			nodePtr_t parent = nullptr;
#if THREAD_SAFE
            std::lock_guard< std::mutex > lck( m_mutex__ );
#endif
			// 查找节点
			while (curr && curr->m_key != key ) {
				parent = curr;

				if (key < curr->m_key) {
					curr = curr->p_left;
				} else {
					curr = curr->p_right;
				}
			}

			if (!curr) { return; }  // 节点不存在，不做任何操作

			if (!curr->p_left && !curr->p_right) { // 节点没有子树
                if (curr == p_root__) {
                    p_root__ = nullptr;
				} else if (curr == parent->p_left) {
					parent->p_left = nullptr;
                    fix_deletion__(parent, nullptr);
				} else {
					parent->p_right = nullptr;
                    fix_deletion__(parent, nullptr);
				}
			} else if (curr->p_left && curr->p_right) {
				stNode* succ = curr->p_right;
				parent = curr;

				while (succ->p_left) {
					parent = succ;
					succ = succ->p_left;
				}

				curr->m_key = succ->m_key;
				curr->m_val = succ->m_val;
			
				if (succ == parent->p_left) {
					parent->p_left = succ->p_right;
				} else {
					parent->p_right = succ->p_right;
				}

                fix_deletion__(parent, succ->p_right);
			} else {
				nodePtr_t child = curr->p_left ? curr->p_left : curr->p_right;

                if (curr == p_root__ ) { // 节点是根节点
                    p_root__ = p_root__->p_left ? p_root__->p_left : p_root__->p_right;
                    p_root__->m_color = BLACK;
				} else if (curr == parent -> p_left) {
					parent->p_left = child;
                    fix_deletion__(parent, child);
				} else {
					parent->p_right = child;
                    fix_deletion__(parent, child);
				}
			}

			delete curr;
		}

	
	private:
		/**
		 * @brief 插入后平衡操作
		 */
        void fix_insertion__( stNode* node) {
			if( node == nullptr ) return;
            nodePtr_t parent = get_parent__( node );
			if( parent == nullptr ) return;
            nodePtr_t grand = get_parent__( parent );
			
			while (parent != nullptr && parent->m_color == RED) {
				if ( parent == grand->p_left) { // 祖父的左孩子
					nodePtr_ty = grand->p_right;
					if (y != nullptr && y->m_color == RED) {
						parent->m_color = BLACK;
						y->m_color = BLACK;
						grand->m_color = RED;
						node = grand;
					} else {
						if (node == parent->p_right) {
							node = parent;
                            rotate_left__( node );
						}
						parent->m_color = BLACK;
						grand->m_color = RED;
                        rotate_right__(grand);
					}
				} else {
					nodePtr_ty = grand->p_left;
					if (y != nullptr && y->m_color == RED) {
						parent->m_color = BLACK;
						y->m_color = BLACK;
						grand->m_color = RED;
						node = grand;
					} else {
						if (node == parent->p_left) {
							node = parent;
                            rotate_right__( node );
						}
						parent->m_color = BLACK;
						grand->m_color = RED;
                        rotate_left__( grand );
					}
				}

                parent = get_parent__( node );
                grand = get_parent__( parent );
			}
            p_root__->m_color = BLACK;
		}

		/**
		 * @brief 删除后平衡操作
		 */
        void fix_deletion__( nodePtr_t parent, nodePtr_t node) {
            while (node != p_root__ && !node->m_color ) {
				if (node == parent->p_left ) {
					stNode* sibling = parent->p_right;

					if (sibling-> m_color ) {
						sibling->m_color = BLACK;
						parent->m_color = RED;
                        rotate_left__(parent);
						sibling = parent->p_right;
					}

					if (!sibling->p_left->m_color && !sibling->p_right->m_color) {
						sibling->m_color = RED;
						node = parent;
                        parent = get_parent__(node);
					} else {
						if (!sibling -> p_right->m_color ) {
							sibling->p_left->m_color = BLACK;
							sibling->m_color = RED;
                            rotate_right__(sibling);
							sibling = parent->p_right;
						}

						sibling->m_color = parent->m_color;
						parent->m_color = BLACK;
						sibling->p_right->m_color = BLACK;
                        rotate_left__(parent);
                        node = p_root__;
					}
				} else {
					nodePtr_t sibling = parent->p_left;

					if (sibling->m_color) {
						sibling->m_color = BLACK;
						parent->m_color = RED;
                        rotate_right__(parent);
						sibling = parent->p_left;
					}

					if (!sibling->p_left->m_color && !sibling->p_right->m_color) {
						sibling->m_color = RED;
						node = parent;
                        parent = get_parent__(node);
					} else {
						if (!sibling->p_left->m_color) {
							sibling->p_right->m_color = BLACK;
							sibling->m_color = RED;
                            rotate_left__(sibling);
							sibling = parent->p_left;
						}

						sibling->m_color = parent->m_color;
						parent->m_color = BLACK;
						sibling->p_left->m_color = BLACK;
                        rotate_right__(parent);
                        node = p_root__;
					}
				}
			}

			if (node) {
				node->m_color = BLACK;
			}
		}
		/**
		 * @brief 取双亲节点
		 */
        nodePtr_t get_parent__( nodePtr_t node) const {
			if (!node) {
				return {};
			}

            if (node == p_root__) {
				return {};
			}

            nodePtr_t curr = p_root__ , parent;

			while (curr && curr != node) {
				parent = curr;

				if (node->m_key < curr->m_key ) {
					curr = curr->p_left;
				} else {
					curr = curr->p_right;
				}
			}

			return parent;
		}
		/**
		 * @brief 取叔叔节点
		 */
		nodePtr_t __get_uncle( nodePtr_t node) const {
            nodePtr_t parent = get_parent__(node);
            nodePtr_t grandparent = get_parent__(parent);

			if (!parent || !grandparent) {
				return nullptr;
			}

			if (parent == grandparent->p_left ) {
				return grandparent->p_right;
			} else {
				return grandparent->p_left;
			}
		}
		/**
		 * @brief 左旋
		 */
        void rotate_left__( nodePtr_t node) {
			nodePtr_t right = node->p_right;
            nodePtr_t parent = get_parent__(node);

            if (node == p_root__ ) {
                p_root__ = node->p_right;
			} else if (node == parent->p_left ) {
				parent->p_left = node->p_right;
			} else {
				parent->p_right = node->p_right;
			}

			node->p_right = right->p_left;
			right->p_left = node;
		}
		/**
		 * @brief 右旋
		 */
        void rotate_right__( nodePtr_t node) {
			nodePtr_t left = node->p_left;
            nodePtr_t parent = get_parent__(node);

            if (node == p_root__ ) {
                p_root__ = node->p_left;
			} else if (node == parent->p_left ) {
				parent->p_left = node->p_left;
			} else {
				parent->p_right = node->p_left;
			}

			node->p_left = left->p_right;
			left->p_right = node;
		}
	};
}
